package com.wy.config.authorization;

import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.ClaimAccessor;
import org.springframework.security.oauth2.core.ClientAuthenticationMethod;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;

import com.wy.util.SecurityHelpers;

import dream.flying.flower.framework.security.constant.ConstSecurity;
import lombok.extern.slf4j.Slf4j;

/**
 * 短信验证码登录认证提供者
 *
 * @author 飞花梦影
 * @date 2024-09-18 22:18:13
 * @git {@link https://github.com/dreamFlyingFlower}
 */
@Slf4j
public class SmsCaptchaGrantAuthenticationProvider implements AuthenticationProvider {

	private OAuth2TokenGenerator<?> tokenGenerator;

	private AuthenticationManager authenticationManager;

	private OAuth2AuthorizationService authorizationService;

	private static final String ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-5.2";

	private static final OAuth2TokenType ID_TOKEN_TOKEN_TYPE = new OAuth2TokenType(OidcParameterNames.ID_TOKEN);

	@Override
	public Authentication authenticate(Authentication authentication) throws AuthenticationException {
		SmsCaptchaGrantAuthenticationToken authenticationToken = (SmsCaptchaGrantAuthenticationToken) authentication;

		// Ensure the client is authenticated
		OAuth2ClientAuthenticationToken clientPrincipal =
				SecurityHelpers.getAuthenticatedClientElseThrowInvalidClient(authenticationToken);
		RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
		// Ensure the client is configured to use this authorization grant type
		if (registeredClient == null || !registeredClient.getAuthorizationGrantTypes()
				.contains(authenticationToken.getAuthorizationGrantType())) {
			throw new OAuth2AuthenticationException(OAuth2ErrorCodes.UNAUTHORIZED_CLIENT);
		}

		// 验证scope
		Set<String> authorizedScopes = getAuthorizedScopes(registeredClient, authenticationToken.getScopes());

		// 进行认证
		Authentication authenticate = getAuthenticatedUser(authenticationToken);

		// 以下内容摘抄自OAuth2AuthorizationCodeAuthenticationProvider
		DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
				.registeredClient(registeredClient)
				.principal(authenticate)
				.authorizationServerContext(AuthorizationServerContextHolder.getContext())
				.authorizedScopes(authorizedScopes)
				.authorizationGrantType(authenticationToken.getAuthorizationGrantType())
				.authorizationGrant(authenticationToken);

		// Initialize the OAuth2Authorization
		OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
				.principalName(clientPrincipal.getName())
				.authorizationGrantType(authenticationToken.getAuthorizationGrantType());

		// ----- Access token -----
		OAuth2TokenContext tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
		OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
		if (generatedAccessToken == null) {
			OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
					"The token generator failed to generate the access token.", ERROR_URI);
			throw new OAuth2AuthenticationException(error);
		}

		if (log.isTraceEnabled()) {
			log.trace("Generated access token");
		}
		OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
				generatedAccessToken.getTokenValue(), generatedAccessToken.getIssuedAt(),
				generatedAccessToken.getExpiresAt(), tokenContext.getAuthorizedScopes());
		if (generatedAccessToken instanceof ClaimAccessor) {
			authorizationBuilder.token(accessToken,
					(metadata) -> metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME,
							((ClaimAccessor) generatedAccessToken).getClaims()));
		} else {
			authorizationBuilder.accessToken(accessToken);
		}
		// ----- Refresh token -----
		OAuth2RefreshToken refreshToken = null;
		if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
		// Do not issue refresh token to public client
				!clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

			tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
			OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
			if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
				OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
						"The token generator failed to generate the refresh token.", ERROR_URI);
				throw new OAuth2AuthenticationException(error);
			}

			if (log.isTraceEnabled()) {
				log.trace("Generated refresh token");
			}
			refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
			authorizationBuilder.refreshToken(refreshToken);
		}

		// ----- ID token -----
		OidcIdToken idToken;
		if (authorizedScopes.contains(OidcScopes.OPENID)) {
			tokenContext = tokenContextBuilder.tokenType(ID_TOKEN_TOKEN_TYPE)
					// ID token customizer may need access to the access token and/or refresh token
					.authorization(authorizationBuilder.build())
					.build();
			// @formatter:on
			OAuth2Token generatedIdToken = this.tokenGenerator.generate(tokenContext);
			if (!(generatedIdToken instanceof Jwt)) {
				OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
						"The token generator failed to generate the ID token.", ERROR_URI);
				throw new OAuth2AuthenticationException(error);
			}

			if (log.isTraceEnabled()) {
				log.trace("Generated id token");
			}

			idToken = new OidcIdToken(generatedIdToken.getTokenValue(), generatedIdToken.getIssuedAt(),
					generatedIdToken.getExpiresAt(), ((Jwt) generatedIdToken).getClaims());
			authorizationBuilder.token(idToken,
					(metadata) -> metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, idToken.getClaims()));
		} else {
			idToken = null;
		}

		OAuth2Authorization authorization = authorizationBuilder.build();

		// Save the OAuth2Authorization
		this.authorizationService.save(authorization);

		Map<String, Object> additionalParameters = new HashMap<>(1);
		if (idToken != null) {
			// 放入idToken
			additionalParameters.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
		}

		return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, accessToken, refreshToken,
				additionalParameters);
	}

	/**
	 * 获取认证过的scope
	 *
	 * @param registeredClient 客户端
	 * @param requestedScopes 请求中的scope
	 * @return 认证过的scope
	 */
	private Set<String> getAuthorizedScopes(RegisteredClient registeredClient, Set<String> requestedScopes) {
		// Default to configured scopes
		Set<String> authorizedScopes = registeredClient.getScopes();
		if (!ObjectUtils.isEmpty(requestedScopes)) {
			Set<String> unauthorizedScopes = requestedScopes.stream()
					.filter(requestedScope -> !registeredClient.getScopes().contains(requestedScope))
					.collect(Collectors.toSet());
			if (!ObjectUtils.isEmpty(unauthorizedScopes)) {
				SecurityHelpers.throwError(OAuth2ErrorCodes.INVALID_REQUEST,
						"OAuth 2.0 Parameter: " + OAuth2ParameterNames.SCOPE, ERROR_URI);
			}

			authorizedScopes = new LinkedHashSet<>(requestedScopes);
		}

		if (log.isTraceEnabled()) {
			log.trace("Validated token request parameters");
		}
		return authorizedScopes;
	}

	/**
	 * 获取认证过的用户信息
	 *
	 * @param authenticationToken converter构建的认证信息，这里是包含手机号与验证码的
	 * @return 认证信息
	 */
	public Authentication getAuthenticatedUser(SmsCaptchaGrantAuthenticationToken authenticationToken) {
		// 获取手机号密码
		Map<String, Object> additionalParameters = authenticationToken.getAdditionalParameters();
		String phone = (String) additionalParameters.get(ConstSecurity.OAUTH_PARAMETER_NAME_PHONE);
		String smsCaptcha = (String) additionalParameters.get(ConstSecurity.OAUTH_PARAMETER_NAME_SMS_CAPTCHA);
		// 构建UsernamePasswordAuthenticationToken通过AbstractUserDetailsAuthenticationProvider及其子类对手机号与验证码进行校验
		// 这里就是我说的短信验证与密码模式区别不大，如果是短信验证模式则在SmsCaptchaLoginAuthenticationProvider中加一个校验，
		// 使框架支持手机号、验证码校验，反之不加就是账号密码登录
		UsernamePasswordAuthenticationToken unauthenticated =
				UsernamePasswordAuthenticationToken.unauthenticated(phone, smsCaptcha);
		Authentication authenticate = null;
		try {
			authenticate = authenticationManager.authenticate(unauthenticated);
		} catch (Exception e) {
			SecurityHelpers.throwError(OAuth2ErrorCodes.INVALID_REQUEST, "认证失败：手机号或验证码错误.", ERROR_URI);
		}
		return authenticate;
	}

	@Override
	public boolean supports(Class<?> authentication) {
		return SmsCaptchaGrantAuthenticationToken.class.isAssignableFrom(authentication);
	}

	public void setTokenGenerator(OAuth2TokenGenerator<?> tokenGenerator) {
		Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
		this.tokenGenerator = tokenGenerator;
	}

	public void setAuthenticationManager(AuthenticationManager authenticationManager) {
		Assert.notNull(authorizationService, "authenticationManager cannot be null");
		this.authenticationManager = authenticationManager;
	}

	public void setAuthorizationService(OAuth2AuthorizationService authorizationService) {
		Assert.notNull(authorizationService, "authorizationService cannot be null");
		this.authorizationService = authorizationService;
	}
}